from channels.generic.websocket import WebsocketConsumer,AsyncWebsocketConsumer
from asgiref.sync import async_to_sync
from channels.exceptions import StopConsumer

from django.http.request import QueryDict
from websocket.ssh import SSH
from asset.models import *
from websocket.tasks import task_exec_add

from utils.custom_log import log_start
logger = log_start('consumer')
import os
import json
import base64
import re

class AnsibleConsumer(WebsocketConsumer):
    def connect(self):
        ''' 当Websocket创建连接时 '''
        self.accept()
    
    def disconnect(self, code):
        ''' 当Websocket发生断开连接时 '''
        pass
    def receive(self, text_data=None, bytes_data=None):
        ''' 当Websocket接收到前端send消息时 '''
        data = json.loads(text_data)
        logger.error(f"接收到前端ws发送的数据:\n {data}")

        query_string = self.scope.get('query_string') ## 获取建立ws后前端入参
        args = QueryDict(query_string=query_string, encoding='utf-8')
        logger.error (f"接收到前端接口请求传参:\n {args}")
        # message = { 'message': 'test'}
        # message = json.dumps(message)
        # self.send(message)
        # self.user = self.scope['user'] # 获取用户
        hosts = args.get('host')
        groups = args.get('group')
        logger.error(f"hosts: {hosts}")
        logger.error(f"groups: {groups}")
        
        result = task_exec_add.delay(
            hosts=hosts,
            groups=groups,
            channel_name=self.channel_name
        )
        data = {'detail': f'异步更新中... ID:{result.id} [ {result.status} ] '}
        logger.info(f"{data}")

    def send_message(self, event):
        #自定义函数，在celery任务中使用channels_layers会用到
        self.send(text_data=json.dumps(event))


        

class SSHConsumer(WebsocketConsumer): 

    message = {'status': 0, 'message': None}

    def connect(self):
        """
        打开 websocket 连接, 通过前端传入的参数尝试连接 ssh 主机
        :return:
        """
        self.accept()
        query_string = self.scope.get('query_string')
        args = QueryDict(query_string=query_string, encoding='utf-8')
        logger.info (f"QueryDict:{args}")
        id = args.get('id') # 获取row id
        HostInstance = Host.objects.get(id=int(id))

        if HostInstance.connect_type == 1:
            host = HostInstance.floatingip
            matchObj=re.match(r'(.*.?\d)?',host)
            if matchObj:
                host =  matchObj.group(1)
        else:
            host = HostInstance.ip
        
        port = HostInstance.ssh_id.ssh_port
        user = HostInstance.ssh_id.ssh_user
        password = HostInstance.ssh_id.password
        private = HostInstance.ssh_id.private

        width = int(args.get('width'))
        height = int(args.get('height'))
        
        ssh_connect_dict = {
            'host': host,
            'user': user,
            'port': port,
            'timeout': 30,
            'width': width,
            'height': height,
            'password': password,
            'ssh_key':private
        }
        
        ## paramiko 建立连接
        try:
            self.ssh = SSH(websocker=self, message=self.message)
            self.ssh.connect(**ssh_connect_dict)
        except Exception as e:
            self.ssh.close("paramiko 建立连接失败")

        
            

    def disconnect(self, close_code):
        try:
            if close_code == 3001:
                logger.info("close_code=3001")
                # self.ssh.close()
                # pass
            else:
                logger.info("self.ssh.close")
        except Exception as  e:
            pass
        finally:
            # 过滤点结果中的颜色字符
            # res = re.sub('(\[\d{2};\d{2}m|\[0m)', '', self.ssh.res)
            logger.info('Websocket disconnect')
 


    def receive(self, text_data=None, bytes_data=None):
        if text_data is None:
            self.ssh.django_bytes_to_ssh(bytes_data)
        else:
            data = json.loads(text_data)
            if type(data) == dict:
                status = data['status']
                if status == 0:
                    data = data['data']
                    logger.info(f"接收前端 Send data: {data}")
                    #  # 调用celery任务
                    # self.result = task_add.delay(data, self.channel_name) # self.channel_name是websocketconsumer本身具有的一个属性
                    self.ssh.django_to_ssh(data)
                    # self.ssh.shell(data)
                else:
                    logger.error(f"调整窗口中cols:{data['cols']},rows:{data['rows']}" )
                    cols = data['cols']
                    rows = data['rows']
                    self.ssh.resize_pty(cols=cols, rows=rows)

    
    def send_message(self, event):
        #自定义函数，在celery任务中使用channels_layers会用到
        self.send(text_data=json.dumps(event))